from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import os
import base64
import json
import requests
from werkzeug.utils import secure_filename
from dotenv import load_dotenv
import tempfile
import shutil
import re
from PIL import Image, ImageDraw
import io
from prompts import get_system_prompt, get_user_prompt, is_polyhedron_detection_mode

# 加载环境变量
load_dotenv()

app = Flask(__name__)
CORS(app)

# 配置
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}
MAX_CONTENT_LENGTH = 10 * 1024 * 1024  # 10MB

# Qwen2.5-VL使用1000x1000标准化坐标空间，无需复杂的resize计算

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 所有resize相关函数已删除，因为Qwen2.5-VL使用1000x1000标准化坐标空间

def convert_coordinates_to_original(coordinates: list, original_width: int, original_height: int) -> list:
    """
    将Qwen2.5-VL的标准化坐标(0-1000)转换为原始图像坐标
    Qwen2.5-VL使用1000x1000标准化坐标空间
    """
    converted_coords = []
    for coord in coordinates:
        if len(coord) >= 4:
            x1, y1, x2, y2 = coord[:4]
            # 从1000x1000标准化空间转换到原始图像坐标
            orig_x1 = int(x1 * original_width / 1000)
            orig_y1 = int(y1 * original_height / 1000)
            orig_x2 = int(x2 * original_width / 1000)
            orig_y2 = int(y2 * original_height / 1000)
            converted_coords.append([orig_x1, orig_y1, orig_x2, orig_y2])
    return converted_coords

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def image_to_base64(image_path):
    """将图片转换为base64编码"""
    try:
        with open(image_path, 'rb') as image_file:
            image_data = image_file.read()
            base64_image = base64.b64encode(image_data).decode('utf-8')
            
            # 根据文件扩展名确定MIME类型
            ext = os.path.splitext(image_path)[1].lower()
            mime_map = {
                '.jpg': 'image/jpeg',
                '.jpeg': 'image/jpeg',
                '.png': 'image/png',
                '.gif': 'image/gif',
                '.webp': 'image/webp'
            }
            mime_type = mime_map.get(ext, 'image/jpeg')
            
            return f"data:{mime_type};base64,{base64_image}"
    except Exception as e:
        print(f"图片转换失败: {e}")
        return None

def call_siliconflow_api(messages, config=None):
    """调用SiliconFlow API"""
    try:
        url = os.getenv('SILICONFLOW_API_URL')
        api_key = os.getenv('SILICONFLOW_API_KEY')
        model_name = os.getenv('MODEL_NAME')
        
        headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
        
        # 使用配置参数或默认值
        temperature = 0.7
        max_tokens = 2000
        
        if config:
            temperature = config.get('temperature', 0.7)
            max_tokens = config.get('maxTokens', 2000)
        
        data = {
            'model': model_name,
            'messages': messages,
            'max_tokens': max_tokens,
            'temperature': temperature
        }
        
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        
        return response.json()
    except Exception as e:
        print(f"API调用失败: {e}")
        raise e

@app.route('/')
def index():
    """主页"""
    return send_from_directory('public', 'index.html')

@app.route('/public/<path:filename>')
def static_files(filename):
    """静态文件服务"""
    return send_from_directory('public', filename)

@app.route('/config')
def config_page():
    """配置页面"""
    return send_from_directory('public', 'config.html')

@app.route('/api/prompts', methods=['GET'])
def get_prompts():
    """获取预设提示词"""
    try:
        from prompts import get_preset_prompts
        return jsonify({
            'success': True,
            'prompts': get_preset_prompts()
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'获取提示词失败: {str(e)}'
        }), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    """处理文本对话"""
    try:
        data = request.get_json()
        message = data.get('message', '').strip()
        history = data.get('history', [])
        config = data.get('config', {})
        
        if not message:
            return jsonify({'error': '消息不能为空'}), 400
        
        # 构建消息历史，添加系统提示词
        messages = []
        
        # 添加系统提示词（如果有配置）
        system_prompt = config.get('systemPrompt', '').strip()
        if system_prompt:
            messages.append({'role': 'system', 'content': system_prompt})
        
        # 添加历史对话和当前消息
        messages.extend(history)
        messages.append({'role': 'user', 'content': message})
        
        # 调用API
        result = call_siliconflow_api(messages, config)
        
        return jsonify({
            'success': True,
            'response': result['choices'][0]['message']['content'],
            'usage': result.get('usage', {})
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'对话失败: {str(e)}'
        }), 500

@app.route('/api/chat-with-image', methods=['POST'])
def chat_with_image():
    """处理图片+文本对话"""
    temp_file_path = None
    
    try:
        message = request.form.get('message', '').strip()
        history_str = request.form.get('history', '[]')
        config_str = request.form.get('config', '{}')
        history = json.loads(history_str) if history_str else []
        config = json.loads(config_str) if config_str else {}
        
        # 获取上传的图片
        image_file = request.files.get('image')
        
        if not message and not image_file:
            return jsonify({'error': '请提供文本消息或图片'}), 400
        
        # 构建消息内容
        content = []
        
        if message:
            content.append({
                'type': 'text',
                'text': message
            })
        
        if image_file and allowed_file(image_file.filename):
            # 保存临时文件
            filename = secure_filename(image_file.filename)
            temp_file_path = os.path.join(app.config['UPLOAD_FOLDER'], 
                                        f"temp_{os.getpid()}_{filename}")
            image_file.save(temp_file_path)
            
            # 转换为base64
            base64_image = image_to_base64(temp_file_path)
            if base64_image:
                content.append({
                    'type': 'image_url',
                    'image_url': {
                        'url': base64_image
                    }
                })
        
        # 构建消息历史，添加系统提示词
        messages = []
        
        # 添加系统提示词（如果有配置）
        system_prompt = config.get('systemPrompt', '').strip()
        if system_prompt:
            messages.append({'role': 'system', 'content': system_prompt})
        
        # 添加历史对话和当前消息
        messages.extend(history)
        messages.append({'role': 'user', 'content': content})
        
        # 调用API
        result = call_siliconflow_api(messages, config)
        response_text = result['choices'][0]['message']['content']
        
        # 检查是否为多面体检测模式，如果是则解析坐标并提供转换信息
        response_data = {
            'success': True,
            'response': response_text,
            'usage': result.get('usage', {})
        }
        
        # 如果是多面体检测模式且有图片，提供坐标转换信息
        if (temp_file_path and 
            is_polyhedron_detection_mode(config.get('systemPrompt', ''))):
            
            # 解析坐标
            coordinates = parse_coordinates_from_response(response_text)
            if coordinates:
                # 获取图像尺寸信息用于前端坐标转换
                image = Image.open(temp_file_path)
                original_width, original_height = image.size
                
                response_data.update({
                    'coordinates': coordinates,
                    'image_info': {
                        'original_width': original_width,
                        'original_height': original_height
                    }
                })
        
        return jsonify(response_data)
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'对话失败: {str(e)}'
        }), 500
    
    finally:
        # 清理临时文件
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
            except:
                pass

def parse_coordinates_from_response(response_text):
    """从模型响应中解析坐标信息"""
    try:
        # 尝试匹配JSON格式的坐标
        json_pattern = r'\{[^{}]*"coordinates"[^{}]*\}'
        json_matches = re.findall(json_pattern, response_text, re.DOTALL)
        
        if json_matches:
            for match in json_matches:
                try:
                    coord_data = json.loads(match)
                    if 'coordinates' in coord_data:
                        return coord_data['coordinates']
                except:
                    continue
        
        # 尝试匹配数组格式的坐标 [x1, y1, x2, y2]
        array_pattern = r'\[(\d+),\s*(\d+),\s*(\d+),\s*(\d+)\]'
        array_matches = re.findall(array_pattern, response_text)
        
        if array_matches:
            return [[int(x) for x in match] for match in array_matches]
        
        # 尝试匹配bbox格式
        bbox_pattern = r'<box>\((\d+),(\d+)\),\((\d+),(\d+)\)</box>'
        bbox_matches = re.findall(bbox_pattern, response_text)
        
        if bbox_matches:
            return [[int(match[0]), int(match[1]), int(match[2]), int(match[3])] for match in bbox_matches]
        
        return None
    except Exception as e:
        print(f"坐标解析失败: {e}")
        return None

def draw_bounding_boxes(image_path, coordinates):
    """在图片上绘制边界框，使用1000x1000标准化坐标转换"""
    try:
        # 打开原图
        image = Image.open(image_path)
        original_width, original_height = image.size
        draw = ImageDraw.Draw(image)
        
        # 转换坐标：从1000x1000标准化空间到原始图像坐标
        converted_coordinates = convert_coordinates_to_original(coordinates, original_width, original_height)
        
        # 输出调试信息
        print(f"[后端坐标转换] 原始尺寸: {original_width}x{original_height}")
        print(f"[后端坐标转换] 坐标系统: 1000x1000标准化空间")
        print(f"[后端坐标转换] AI返回坐标: {coordinates}")
        print(f"[后端坐标转换] 转换后坐标: {converted_coordinates}")
        
        # 绘制每个边界框
        for coord in converted_coordinates:
            if len(coord) >= 4:
                x1, y1, x2, y2 = coord[:4]
                # 确保坐标在图像范围内
                x1 = max(0, min(x1, original_width))
                y1 = max(0, min(y1, original_height))
                x2 = max(0, min(x2, original_width))
                y2 = max(0, min(y2, original_height))
                
                # 绘制红色边界框，线宽为3
                draw.rectangle([x1, y1, x2, y2], outline='red', width=3)
        
        # 保存到内存中
        img_buffer = io.BytesIO()
        image.save(img_buffer, format='PNG')
        img_buffer.seek(0)
        
        # 转换为base64
        img_data = img_buffer.getvalue()
        base64_image = base64.b64encode(img_data).decode('utf-8')
        return f"data:image/png;base64,{base64_image}"
        
    except Exception as e:
        print(f"绘制边界框失败: {e}")
        return None

@app.route('/api/detect-polyhedron', methods=['POST'])
def detect_polyhedron():
    """检测图片中的多面体并返回坐标"""
    temp_file_path = None
    
    try:
        # 获取参数
        config_str = request.form.get('config', '{}')
        config = json.loads(config_str) if config_str else {}
        
        # Qwen2.5-VL使用1000x1000标准化坐标，不需要处理resize参数
        
        # 获取上传的图片
        image_file = request.files.get('image')
        
        if not image_file or not allowed_file(image_file.filename):
            return jsonify({'error': '请上传有效的图片文件'}), 400
        
        # 保存临时文件
        filename = secure_filename(image_file.filename)
        temp_file_path = os.path.join(app.config['UPLOAD_FOLDER'], 
                                    f"polyhedron_{os.getpid()}_{filename}")
        image_file.save(temp_file_path)
        
        # 转换为base64
        base64_image = image_to_base64(temp_file_path)
        if not base64_image:
            return jsonify({'error': '图片处理失败'}), 500
        
        # 获取多面体检测提示词
        polyhedron_prompt = get_user_prompt("polyhedron_detection")
        
        # 构建消息
        content = [
            {
                'type': 'text',
                'text': polyhedron_prompt
            },
            {
                'type': 'image_url',
                'image_url': {
                    'url': base64_image
                }
            }
        ]
        
        messages = [
            {'role': 'system', 'content': get_system_prompt("vision_assistant")},
            {'role': 'user', 'content': content}
        ]
        
        # 调用API
        result = call_siliconflow_api(messages, config)
        response_text = result['choices'][0]['message']['content']
        
        # 解析坐标
        coordinates = parse_coordinates_from_response(response_text)
        
        # 如果找到坐标，绘制边界框
        annotated_image = None
        if coordinates:
            annotated_image = draw_bounding_boxes(temp_file_path, coordinates)
        
        return jsonify({
            'success': True,
            'response': response_text,
            'coordinates': coordinates,
            'annotated_image': annotated_image,
            'usage': result.get('usage', {})
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'多面体检测失败: {str(e)}'
        }), 500
    
    finally:
        # 清理临时文件
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
            except:
                pass

if __name__ == '__main__':
    port = int(os.getenv('PORT', 3000))
    print(f"服务器运行在 http://localhost:{port}")
    print(f"也可以通过IP地址访问: http://0.0.0.0:{port}")
    print("请确保在.env文件中配置了正确的API密钥")
    
    app.run(host='0.0.0.0', port=port, debug=True) 